from bs4 import BeautifulSoup
import os
import re
from collections import Counter

# 加载 TEI 文件
tei_file_path = './oth.xml'
with open(tei_file_path, 'r', encoding='utf-8') as file:
    soup = BeautifulSoup(file, 'xml')

# 提取标题
title_tag = soup.find('title')

# 提取 body 中的所有文本
body_text = []
for w_tag in soup.find_all('w'):
    body_text.append(w_tag.get_text())
body_text_str = ' '.join(body_text)

# 创建一个新的纯文本文件，文件名根据提取的标题来构建
title = title_tag.get_text() if title_tag else 'Othello'
title_cleaned = re.sub(r'\W+', '_', title.lower()).strip('_')  # 清理标题以用于文件名

# 创建新的纯文本文件
output_file_path = f'{title_cleaned}.txt'
with open(output_file_path, 'w', encoding='utf-8') as file:
    file.write(body_text_str)

# 提取 Iago 台词中的所有文本并计算单词出现次数
iago_speeches = soup.find_all('speaker', string='Iago')
iago_text = []
for speech in iago_speeches:
    # 找到与当前 speaker 标签相邻的 w 标签
    sibling = speech.find_next_sibling('w')
    while sibling and sibling.find_parent('speaker') is None:  # 直到遇到下一个 speaker 标签或没有更多 w 标签
        iago_text.append(sibling.get_text())
        sibling = sibling.find_next_sibling('w')

iago_text_str = ' '.join(iago_text)
word_counts = Counter(re.findall(r'\b\w+\b', iago_text_str))

# 创建存储单词计数的文件
word_counts_file_path = 'iago_word_counts.txt'
with open(word_counts_file_path, 'w', encoding='utf-8') as file:
    for word, count in word_counts.items():
        file.write(f'{word} {count}\n')

# 定义函数来清理单词
def clean_word(word):
    return re.sub(r'[^\w\s]', '', word).strip().lower()

# 计算每幕的单词总数并找出单词数最多的那一幕
acts = {}
current_act = None
for line in soup.find_all('l'):  # 假设每行标签包含幕和场景信息以及台词
    # 解析幕信息（这里需要根据实际的 TEI 文件结构进行调整）
    match = re.match(r'^(\d+\.)', line.get('n'))  # 假设行号以幕号开头，如 "1." 表示第一幕
    if match:
        current_act = int(match.group(1).strip('.'))
        acts[current_act] = []
    # 提取并清理单词
    words = line.find_all('w')
    cleaned_words = [clean_word(w.get_text()) for w in words if clean_word(w.get_text())]
    acts[current_act].extend(cleaned_words)

# 找出单词数最多的那一幕
max_words_act = max(acts, key=lambda act: len(acts[act]))
print(f'The act with the most words is Act {max_words_act} with {len(acts[max_words_act])} words.')